package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
)

type OllamaRequest struct {
	Model  string `json:"model"`
	Prompt string `json:"prompt"`
}

type OllamaResponse struct {
	Response string `json:"response"`
}

func main() {
	http.HandleFunc("/ollama", func(w http.ResponseWriter, r *http.Request) {
		var req OllamaRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		// 调用Ollama API
		ollamaReq, _ := json.Marshal(req)
		resp, err := http.Post("http://localhost:11434/api/generate", "application/json", bytes.NewBuffer(ollamaReq))
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		defer resp.Body.Close()

		var ollamaResp OllamaResponse
		json.NewDecoder(resp.Body).Decode(&ollamaResp)

		json.NewEncoder(w).Encode(ollamaResp)
	})

	fmt.Println("Server running on :8080")
	http.ListenAndServe(":8080", nil)
}
